# Copyright 2018-2023 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Unit tests for the ``two_particle`` function of FermionOperator
in openfermion.
"""
import os

import numpy as np
import pytest

from pennylane import qchem

ref_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "test_ref_files")

v_op_1 = {
    (): 0.0,
    ((0, 1), (0, 1), (0, 0), (0, 0)): 0.3411947665760211,
    ((0, 1), (1, 1), (1, 0), (0, 0)): 0.3411947665760211,
    ((1, 1), (0, 1), (0, 0), (1, 0)): 0.3411947665760211,
    ((1, 1), (1, 1), (1, 0), (1, 0)): 0.3411947665760211,
    ((0, 1), (0, 1), (2, 0), (2, 0)): 0.08950028803070323,
    ((0, 1), (1, 1), (3, 0), (2, 0)): 0.08950028803070323,
    ((1, 1), (0, 1), (2, 0), (3, 0)): 0.08950028803070323,
    ((1, 1), (1, 1), (3, 0), (3, 0)): 0.08950028803070323,
    ((0, 1), (2, 1), (0, 0), (2, 0)): 0.08950028803070323,
    ((0, 1), (3, 1), (1, 0), (2, 0)): 0.08950028803070323,
    ((1, 1), (2, 1), (0, 0), (3, 0)): 0.08950028803070323,
    ((1, 1), (3, 1), (1, 0), (3, 0)): 0.08950028803070323,
    ((0, 1), (2, 1), (2, 0), (0, 0)): 0.3353663891543792,
    ((0, 1), (3, 1), (3, 0), (0, 0)): 0.3353663891543792,
    ((1, 1), (2, 1), (2, 0), (1, 0)): 0.3353663891543792,
    ((1, 1), (3, 1), (3, 0), (1, 0)): 0.3353663891543792,
    ((2, 1), (0, 1), (0, 0), (2, 0)): 0.3353663891543792,
    ((2, 1), (1, 1), (1, 0), (2, 0)): 0.3353663891543792,
    ((3, 1), (0, 1), (0, 0), (3, 0)): 0.3353663891543792,
    ((3, 1), (1, 1), (1, 0), (3, 0)): 0.3353663891543792,
    ((2, 1), (0, 1), (2, 0), (0, 0)): 0.08950028803070323,
    ((2, 1), (1, 1), (3, 0), (0, 0)): 0.08950028803070323,
    ((3, 1), (0, 1), (2, 0), (1, 0)): 0.08950028803070323,
    ((3, 1), (1, 1), (3, 0), (1, 0)): 0.08950028803070323,
    ((2, 1), (2, 1), (0, 0), (0, 0)): 0.08950028803070323,
    ((2, 1), (3, 1), (1, 0), (0, 0)): 0.08950028803070323,
    ((3, 1), (2, 1), (0, 0), (1, 0)): 0.08950028803070323,
    ((3, 1), (3, 1), (1, 0), (1, 0)): 0.08950028803070323,
    ((2, 1), (2, 1), (2, 0), (2, 0)): 0.352552816086392,
    ((2, 1), (3, 1), (3, 0), (2, 0)): 0.352552816086392,
    ((3, 1), (2, 1), (2, 0), (3, 0)): 0.352552816086392,
    ((3, 1), (3, 1), (3, 0), (3, 0)): 0.352552816086392,
}

v_op_2 = {
    (): 0.6823895331520422,
    ((0, 1), (0, 0)): 1.1624649805561105,
    ((1, 1), (1, 0)): 1.1624649805561105,
    ((0, 1), (0, 1), (0, 0), (0, 0)): 0.352552816086392,
    ((0, 1), (1, 1), (1, 0), (0, 0)): 0.352552816086392,
    ((1, 1), (0, 1), (0, 0), (1, 0)): 0.352552816086392,
    ((1, 1), (1, 1), (1, 0), (1, 0)): 0.352552816086392,
}

v_op_3 = {
    (): 1.6585666870874103,
    ((0, 1), (0, 0)): 0.7200645027092623,
    ((1, 1), (1, 0)): 0.7200645027092623,
    ((0, 1), (2, 0)): 0.01568677271778423,
    ((1, 1), (3, 0)): 0.01568677271778423,
    ((2, 1), (0, 0)): 0.01568677271778432,
    ((3, 1), (1, 0)): 0.01568677271778432,
    ((2, 1), (2, 0)): 0.7696051973390092,
    ((3, 1), (3, 0)): 0.7696051973390092,
    ((0, 1), (0, 1), (0, 0), (0, 0)): 0.24365548437056841,
    ((0, 1), (1, 1), (1, 0), (0, 0)): 0.24365548437056841,
    ((1, 1), (0, 1), (0, 0), (1, 0)): 0.24365548437056841,
    ((1, 1), (1, 1), (1, 0), (1, 0)): 0.24365548437056841,
    ((0, 1), (0, 1), (0, 0), (2, 0)): -0.02428978770367371,
    ((0, 1), (1, 1), (1, 0), (2, 0)): -0.02428978770367371,
    ((1, 1), (0, 1), (0, 0), (3, 0)): -0.02428978770367371,
    ((1, 1), (1, 1), (1, 0), (3, 0)): -0.02428978770367371,
    ((0, 1), (0, 1), (2, 0), (0, 0)): -0.024289787703673717,
    ((0, 1), (1, 1), (3, 0), (0, 0)): -0.024289787703673717,
    ((1, 1), (0, 1), (2, 0), (1, 0)): -0.024289787703673717,
    ((1, 1), (1, 1), (3, 0), (1, 0)): -0.024289787703673717,
    ((0, 1), (0, 1), (2, 0), (2, 0)): 0.006531987971873417,
    ((0, 1), (1, 1), (3, 0), (2, 0)): 0.006531987971873417,
    ((1, 1), (0, 1), (2, 0), (3, 0)): 0.006531987971873417,
    ((1, 1), (1, 1), (3, 0), (3, 0)): 0.006531987971873417,
    ((0, 1), (2, 1), (0, 0), (0, 0)): -0.02428978770367371,
    ((0, 1), (3, 1), (1, 0), (0, 0)): -0.02428978770367371,
    ((1, 1), (2, 1), (0, 0), (1, 0)): -0.02428978770367371,
    ((1, 1), (3, 1), (1, 0), (1, 0)): -0.02428978770367371,
    ((0, 1), (2, 1), (0, 0), (2, 0)): 0.006531987971873418,
    ((0, 1), (3, 1), (1, 0), (2, 0)): 0.006531987971873418,
    ((1, 1), (2, 1), (0, 0), (3, 0)): 0.006531987971873418,
    ((1, 1), (3, 1), (1, 0), (3, 0)): 0.006531987971873418,
    ((0, 1), (2, 1), (2, 0), (0, 0)): 0.11180501845367194,
    ((0, 1), (3, 1), (3, 0), (0, 0)): 0.11180501845367194,
    ((1, 1), (2, 1), (2, 0), (1, 0)): 0.11180501845367194,
    ((1, 1), (3, 1), (3, 0), (1, 0)): 0.11180501845367194,
    ((0, 1), (2, 1), (2, 0), (2, 0)): 0.0037420836721733935,
    ((0, 1), (3, 1), (3, 0), (2, 0)): 0.0037420836721733935,
    ((1, 1), (2, 1), (2, 0), (3, 0)): 0.0037420836721733935,
    ((1, 1), (3, 1), (3, 0), (3, 0)): 0.0037420836721733935,
    ((2, 1), (0, 1), (0, 0), (0, 0)): -0.02428978770367371,
    ((2, 1), (1, 1), (1, 0), (0, 0)): -0.02428978770367371,
    ((3, 1), (0, 1), (0, 0), (1, 0)): -0.02428978770367371,
    ((3, 1), (1, 1), (1, 0), (1, 0)): -0.02428978770367371,
    ((2, 1), (0, 1), (0, 0), (2, 0)): 0.11180501845367194,
    ((2, 1), (1, 1), (1, 0), (2, 0)): 0.11180501845367194,
    ((3, 1), (0, 1), (0, 0), (3, 0)): 0.11180501845367194,
    ((3, 1), (1, 1), (1, 0), (3, 0)): 0.11180501845367194,
    ((2, 1), (0, 1), (2, 0), (0, 0)): 0.006531987971873418,
    ((2, 1), (1, 1), (3, 0), (0, 0)): 0.006531987971873418,
    ((3, 1), (0, 1), (2, 0), (1, 0)): 0.006531987971873418,
    ((3, 1), (1, 1), (3, 0), (1, 0)): 0.006531987971873418,
    ((2, 1), (0, 1), (2, 0), (2, 0)): 0.003742083672173376,
    ((2, 1), (1, 1), (3, 0), (2, 0)): 0.003742083672173376,
    ((3, 1), (0, 1), (2, 0), (3, 0)): 0.003742083672173376,
    ((3, 1), (1, 1), (3, 0), (3, 0)): 0.003742083672173376,
    ((2, 1), (2, 1), (0, 0), (0, 0)): 0.006531987971873423,
    ((2, 1), (3, 1), (1, 0), (0, 0)): 0.006531987971873423,
    ((3, 1), (2, 1), (0, 0), (1, 0)): 0.006531987971873423,
    ((3, 1), (3, 1), (1, 0), (1, 0)): 0.006531987971873423,
    ((2, 1), (2, 1), (0, 0), (2, 0)): 0.0037420836721733645,
    ((2, 1), (3, 1), (1, 0), (2, 0)): 0.0037420836721733645,
    ((3, 1), (2, 1), (0, 0), (3, 0)): 0.0037420836721733645,
    ((3, 1), (3, 1), (1, 0), (3, 0)): 0.0037420836721733645,
    ((2, 1), (2, 1), (2, 0), (0, 0)): 0.0037420836721733727,
    ((2, 1), (3, 1), (3, 0), (0, 0)): 0.0037420836721733727,
    ((3, 1), (2, 1), (2, 0), (1, 0)): 0.0037420836721733727,
    ((3, 1), (3, 1), (3, 0), (1, 0)): 0.0037420836721733727,
    ((2, 1), (2, 1), (2, 0), (2, 0)): 0.16894113834436625,
    ((2, 1), (3, 1), (3, 0), (2, 0)): 0.16894113834436625,
    ((3, 1), (2, 1), (2, 0), (3, 0)): 0.16894113834436625,
    ((3, 1), (3, 1), (3, 0), (3, 0)): 0.16894113834436625,
}

v_op_4 = {
    (): 15.477079668766912,
    ((0, 1), (0, 0)): 4.348718720149925,
    ((1, 1), (1, 0)): 4.348718720149925,
    ((2, 1), (2, 0)): 4.7839298822632825,
    ((3, 1), (3, 0)): 4.7839298822632825,
    ((4, 1), (4, 0)): 3.9720173194427786,
    ((5, 1), (5, 0)): 3.9720173194427786,
    ((0, 1), (0, 1), (0, 0), (0, 0)): 0.3913181430816212,
    ((0, 1), (1, 1), (1, 0), (0, 0)): 0.3913181430816212,
    ((1, 1), (0, 1), (0, 0), (1, 0)): 0.3913181430816212,
    ((1, 1), (1, 1), (1, 0), (1, 0)): 0.3913181430816212,
    ((0, 1), (0, 1), (2, 0), (2, 0)): 0.02795519311163234,
    ((0, 1), (1, 1), (3, 0), (2, 0)): 0.02795519311163234,
    ((1, 1), (0, 1), (2, 0), (3, 0)): 0.02795519311163234,
    ((1, 1), (1, 1), (3, 0), (3, 0)): 0.02795519311163234,
    ((0, 1), (0, 1), (4, 0), (4, 0)): 0.03448849005311543,
    ((0, 1), (1, 1), (5, 0), (4, 0)): 0.03448849005311543,
    ((1, 1), (0, 1), (4, 0), (5, 0)): 0.03448849005311543,
    ((1, 1), (1, 1), (5, 0), (5, 0)): 0.03448849005311543,
    ((0, 1), (2, 1), (0, 0), (2, 0)): 0.02795519311163234,
    ((0, 1), (3, 1), (1, 0), (2, 0)): 0.02795519311163234,
    ((1, 1), (2, 1), (0, 0), (3, 0)): 0.02795519311163234,
    ((1, 1), (3, 1), (1, 0), (3, 0)): 0.02795519311163234,
    ((0, 1), (2, 1), (2, 0), (0, 0)): 0.36443868319762635,
    ((0, 1), (3, 1), (3, 0), (0, 0)): 0.36443868319762635,
    ((1, 1), (2, 1), (2, 0), (1, 0)): 0.36443868319762635,
    ((1, 1), (3, 1), (3, 0), (1, 0)): 0.36443868319762635,
    ((0, 1), (4, 1), (0, 0), (4, 0)): 0.03448849005311544,
    ((0, 1), (5, 1), (1, 0), (4, 0)): 0.03448849005311544,
    ((1, 1), (4, 1), (0, 0), (5, 0)): 0.03448849005311544,
    ((1, 1), (5, 1), (1, 0), (5, 0)): 0.03448849005311544,
    ((0, 1), (4, 1), (4, 0), (0, 0)): 0.3041058444913643,
    ((0, 1), (5, 1), (5, 0), (0, 0)): 0.3041058444913643,
    ((1, 1), (4, 1), (4, 0), (1, 0)): 0.3041058444913643,
    ((1, 1), (5, 1), (5, 0), (1, 0)): 0.3041058444913643,
    ((2, 1), (0, 1), (0, 0), (2, 0)): 0.36443868319762635,
    ((2, 1), (1, 1), (1, 0), (2, 0)): 0.36443868319762635,
    ((3, 1), (0, 1), (0, 0), (3, 0)): 0.36443868319762635,
    ((3, 1), (1, 1), (1, 0), (3, 0)): 0.36443868319762635,
    ((2, 1), (0, 1), (2, 0), (0, 0)): 0.02795519311163234,
    ((2, 1), (1, 1), (3, 0), (0, 0)): 0.02795519311163234,
    ((3, 1), (0, 1), (2, 0), (1, 0)): 0.02795519311163234,
    ((3, 1), (1, 1), (3, 0), (1, 0)): 0.02795519311163234,
    ((2, 1), (2, 1), (0, 0), (0, 0)): 0.02795519311163234,
    ((2, 1), (3, 1), (1, 0), (0, 0)): 0.02795519311163234,
    ((3, 1), (2, 1), (0, 0), (1, 0)): 0.02795519311163234,
    ((3, 1), (3, 1), (1, 0), (1, 0)): 0.02795519311163234,
    ((2, 1), (2, 1), (2, 0), (2, 0)): 0.44007954668752236,
    ((2, 1), (3, 1), (3, 0), (2, 0)): 0.44007954668752236,
    ((3, 1), (2, 1), (2, 0), (3, 0)): 0.44007954668752236,
    ((3, 1), (3, 1), (3, 0), (3, 0)): 0.44007954668752236,
    ((2, 1), (2, 1), (4, 0), (4, 0)): 0.012175350836228654,
    ((2, 1), (3, 1), (5, 0), (4, 0)): 0.012175350836228654,
    ((3, 1), (2, 1), (4, 0), (5, 0)): 0.012175350836228654,
    ((3, 1), (3, 1), (5, 0), (5, 0)): 0.012175350836228654,
    ((2, 1), (4, 1), (2, 0), (4, 0)): 0.012175350836228654,
    ((2, 1), (5, 1), (3, 0), (4, 0)): 0.012175350836228654,
    ((3, 1), (4, 1), (2, 0), (5, 0)): 0.012175350836228654,
    ((3, 1), (5, 1), (3, 0), (5, 0)): 0.012175350836228654,
    ((2, 1), (4, 1), (4, 0), (2, 0)): 0.3124924051344332,
    ((2, 1), (5, 1), (5, 0), (2, 0)): 0.3124924051344332,
    ((3, 1), (4, 1), (4, 0), (3, 0)): 0.3124924051344332,
    ((3, 1), (5, 1), (5, 0), (3, 0)): 0.3124924051344332,
    ((4, 1), (0, 1), (0, 0), (4, 0)): 0.3041058444913643,
    ((4, 1), (1, 1), (1, 0), (4, 0)): 0.3041058444913643,
    ((5, 1), (0, 1), (0, 0), (5, 0)): 0.3041058444913643,
    ((5, 1), (1, 1), (1, 0), (5, 0)): 0.3041058444913643,
    ((4, 1), (0, 1), (4, 0), (0, 0)): 0.034488490053115425,
    ((4, 1), (1, 1), (5, 0), (0, 0)): 0.034488490053115425,
    ((5, 1), (0, 1), (4, 0), (1, 0)): 0.034488490053115425,
    ((5, 1), (1, 1), (5, 0), (1, 0)): 0.034488490053115425,
    ((4, 1), (2, 1), (2, 0), (4, 0)): 0.3124924051344332,
    ((4, 1), (3, 1), (3, 0), (4, 0)): 0.3124924051344332,
    ((5, 1), (2, 1), (2, 0), (5, 0)): 0.3124924051344332,
    ((5, 1), (3, 1), (3, 0), (5, 0)): 0.3124924051344332,
    ((4, 1), (2, 1), (4, 0), (2, 0)): 0.012175350836228654,
    ((4, 1), (3, 1), (5, 0), (2, 0)): 0.012175350836228654,
    ((5, 1), (2, 1), (4, 0), (3, 0)): 0.012175350836228654,
    ((5, 1), (3, 1), (5, 0), (3, 0)): 0.012175350836228654,
    ((4, 1), (4, 1), (0, 0), (0, 0)): 0.03448849005311543,
    ((4, 1), (5, 1), (1, 0), (0, 0)): 0.03448849005311543,
    ((5, 1), (4, 1), (0, 0), (1, 0)): 0.03448849005311543,
    ((5, 1), (5, 1), (1, 0), (1, 0)): 0.03448849005311543,
    ((4, 1), (4, 1), (2, 0), (2, 0)): 0.012175350836228654,
    ((4, 1), (5, 1), (3, 0), (2, 0)): 0.012175350836228654,
    ((5, 1), (4, 1), (2, 0), (3, 0)): 0.012175350836228654,
    ((5, 1), (5, 1), (3, 0), (3, 0)): 0.012175350836228654,
    ((4, 1), (4, 1), (4, 0), (4, 0)): 0.3097576511847101,
    ((4, 1), (5, 1), (5, 0), (4, 0)): 0.3097576511847101,
    ((5, 1), (4, 1), (4, 0), (5, 0)): 0.3097576511847101,
    ((5, 1), (5, 1), (5, 0), (5, 0)): 0.3097576511847101,
}


@pytest.mark.parametrize(
    ("name", "core", "active", "v_op_exp"),
    [
        ("h2_pyscf", None, None, v_op_1),
        ("h2_pyscf", [0], None, v_op_2),
        ("h2_pyscf", None, [0, 1], v_op_1),
        ("lih", [0], [1, 2], v_op_3),
        ("h2o_psi4", [0, 1, 2], [3, 4, 6], v_op_4),
    ],
)
@pytest.mark.usefixtures("skip_if_no_openfermion_support")
def test_table_two_particle(name, core, active, v_op_exp):
    r"""Test the FermionOperator built by the function `two_particle` of the `obs` module"""
    import openfermion

    hf_data = openfermion.MolecularData(filename=os.path.join(ref_dir, name))

    v_op = qchem.two_particle(hf_data.two_body_integrals, core=core, active=active)

    assert v_op.terms == v_op_exp


v_me_1D = np.array([1, 2, 3, 4])
v_me_4D = np.full((2, 2, 2, 2), 0.5)


@pytest.mark.parametrize(
    ("v_me", "core", "active", "msg_match"),
    [
        (v_me_1D, [0], None, "'matrix_elements' must be a 4D array"),
        (v_me_4D, [-1, 0, 1, 2], None, "Indices of core orbitals must be between 0 and"),
        (v_me_4D, [0, 1, 2, 3], None, "Indices of core orbitals must be between 0 and"),
        (v_me_4D, None, [-1, 0], "Indices of active orbitals must be between 0 and"),
        (v_me_4D, None, [2, 6], "Indices of active orbitals must be between 0 and"),
    ],
)
@pytest.mark.usefixtures("skip_if_no_openfermion_support")
def test_exceptions_two_particle(v_me, core, active, msg_match):
    """Test that the function `'two_particle'` throws an exception
    if the dimension of the matrix elements array is not a 4D array or
    if the indices of core and/or active orbitals are out of range."""

    with pytest.raises(ValueError, match=msg_match):
        qchem.two_particle(v_me, core=core, active=active)
